import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

from config import *
from models import DM_SVM, DM_NN, DM_ELM, oneD_Net_train, oneD_Net_predict, oneD_Net
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


def get_dataset(feature_data):
    x = feature_data[:, in_feature_index]
    y = feature_data[:, out_feature_index]
    # shuffle
    state = np.random.get_state()
    np.random.shuffle(x)
    np.random.set_state(state)
    np.random.shuffle(y)
    # standardization
    scaler = StandardScaler()
    x = scaler.fit_transform(x)
    # split
    return train_test_split(x, y, test_size=0.15)


if __name__ == "__main__":
    feature_data = np.load(washed_feature_data_file_path)
    x_train, x_test, y_train, y_test = get_dataset(feature_data)

    # Task 3
    # ELM (Extreme Learning Machine)
    dm_elm = DM_ELM(x_train, x_test, y_train, y_test)
    logger.info("ELM training")
    dm_elm = DM_ELM(x_train, x_test, y_train, y_test)
    dm_elm.train()
    logger.info("ELM predicting")
    DM_ELM_acc, DM_ELM_pred = dm_elm.test()

    # SVM (Support Vector Machine)
    dm_svm = DM_SVM(x_train, x_test, y_train, y_test)
    logger.info("SVM training")
    dm_svm.train()
    logger.info("SVM predicting")
    DM_SVM_acc, DM_SVM_pred = dm_svm.predict()

    # NN (Neural Network)
    dm_nn = DM_NN(x_train, x_test, y_train, y_test)
    logger.info("NN training")
    dm_nn.train(BATCH_SIZE, EPOCHS, LR)
    logger.info("NN evaluating")
    DM_NN_acc, DM_NN_pred = dm_nn.evaluate()
    # dm_nn.save_model(nn_model_checkpoint_file_path)

    dm_oneD_net = oneD_Net(
        input_channel=len(in_feature_name_list), out_channel=len(out_feature_name_list)
    )
    logger.info("oneD_Net training")
    oneD_Net_train(dm_oneD_net, x_train, y_train, 64, 500, 0.05)
    logger.info("oneD_Net evaluating")
    DM_oneD_Net_acc, DM_oneD_Net_pred = oneD_Net_predict(dm_oneD_net, x_test, y_test)

    DM_ELM_acc = pd.DataFrame(DM_ELM_acc.reshape(1, DM_ELM_acc.shape[0]))
    DM_ELM_pred = pd.DataFrame(DM_ELM_pred)

    DM_SVM_acc = pd.DataFrame(DM_SVM_acc.reshape(1, DM_SVM_acc.shape[0]))
    DM_SVM_pred = pd.DataFrame(DM_SVM_pred)

    DM_NN_acc = pd.DataFrame(DM_NN_acc.reshape(1, DM_NN_acc.shape[0]))
    DM_NN_pred = pd.DataFrame(DM_NN_pred)

    DM_oneD_Net_acc = pd.DataFrame(DM_oneD_Net_acc.reshape(1, DM_oneD_Net_acc.shape[0]))
    DM_oneD_Net_pred = pd.DataFrame(DM_oneD_Net_pred)

    # excel save
    with pd.ExcelWriter(
        excel_washed_result_file_path,
        mode="a",
        engine="openpyxl",
        if_sheet_exists="replace",
    ) as writer:
        DM_ELM_acc.to_excel(
            writer,
            sheet_name="任务3_ELM_精确度",
            header=out_feature_name_list,
            index_label="精确度",
            float_format="%.5f",
        )

        DM_ELM_pred.to_excel(
            writer,
            sheet_name="任务3_ELM_预测值",
            header=out_feature_name_list,
            index_label="预测值",
            float_format="%.5f",
        )

        DM_SVM_acc.to_excel(
            writer,
            sheet_name="任务3_SVM_精确度",
            header=out_feature_name_list,
            index_label="精确度",
            float_format="%.5f",
        )

        DM_SVM_pred.to_excel(
            writer,
            sheet_name="任务3_SVM_预测值",
            header=out_feature_name_list,
            index_label="预测值",
            float_format="%.5f",
        )

        DM_NN_acc.to_excel(
            writer,
            sheet_name="任务3_NN_精确度",
            header=out_feature_name_list,
            index_label="精确度",
            float_format="%.5f",
        )
        DM_NN_pred.to_excel(
            writer,
            sheet_name="任务3_NN_预测值",
            header=out_feature_name_list,
            index_label="预测值",
            float_format="%.5f",
        )
        DM_oneD_Net_acc.to_excel(
            writer,
            sheet_name="任务3_1d卷积_精确度",
            header=out_feature_name_list,
            index_label="精确度",
            float_format="%.5f",
        )
        DM_oneD_Net_pred.to_excel(
            writer,
            sheet_name="任务3_1d卷积_预测值",
            header=out_feature_name_list,
            index_label="预测值",
            float_format="%.5f",
        )
    logger.info("Done")
